#include "trainingcontrol.h"
#include "ui_trainingcontrol.h"

TrainingControl::TrainingControl(QWidget *parent) :
    QDialog(parent),
    ui(new Ui::TrainingControl), gt(new GermTools),
    runNNFlag(false)

{
    ui->setupUi(this);

    ui->gridComboBox->addItem("16");
    ui->gridComboBox->addItem("32");
    ui->gridComboBox->addItem("64");
    ui->gridComboBox->addItem("128");

    rubberBand = new QRubberBand(QRubberBand::Rectangle, this);
    origin = QPoint(ui->imgLabel->x(), ui->imgLabel->y());

    connect(ui->manualButton, SIGNAL(clicked()), this, SLOT(manualD()));
    connect(ui->autoButton, SIGNAL(clicked()), this, SLOT(autoD()));
    connect(ui->resetButton, SIGNAL(clicked()), this, SLOT(reset()));
    connect(ui->cancelButton, SIGNAL(clicked()), this, SLOT(cancel()));    
    connect(ui->gridComboBox, SIGNAL(currentIndexChanged(QString)), this, SLOT(setGrid(QString)));
}

TrainingControl::~TrainingControl()
{
    delete ui;
    delete gt;
}

/*
 *  Once NN is trained, use it to test unknow images
 *  and set 'runNNFlag' flag to be true to indicate neural
 *  network has been trained.
 */
void TrainingControl::startRunning()
{
    runNNFlag = true;
}


/*
 *  Add current image in image loader to the trainingcontrol
 *  and wait for user to manually draw ROI on the image
 *
 */
void TrainingControl::addFile(const QString& file)
{   
    Mat src = imread(file.toStdString(), -1);
    if(!src.data) return;

    img = src.clone();
    resultImg = src.clone();
    dst = img.clone();

    // display original image in the training control
    display(src);
    // display original image in displaying icon
    displayIcon(src, Rect());
}

/*
 *  Display image in label and initiates pixmap of src
 *
 */
void TrainingControl::display(Mat &src)
{
    if(!src.data) return;
    Mat rgb;
    cvtColor(src, rgb, CV_BGR2RGB);

    QImage qTemp = QImage(rgb.data, rgb.cols, rgb.rows, rgb.step, QImage::Format_RGB888);

    if(!qTemp.isNull()) qTemp = qTemp.scaled(ui->imgLabel->size(), Qt::KeepAspectRatio);
    else return;

    QPixmap pix = QPixmap::fromImage(qTemp, 0);

    // Set pix to be main pixmap
    ui->imgLabel->setPixmap(pix);
}

void TrainingControl::displayIcon(Mat &src, Rect &roi)
{
    double r1 = ui->iconLabel->width() / ui->imgLabel->width();
    double r2 = ui->iconLabel->height() / ui->imgLabel->height();

    int x = (int) (roi.x * r1);
    int y = (int) (roi.y * r2);
    int w = (int) (roi.width * r1);
    int h = (int) (roi.height * r2);

    if(!src.data) return;
    Mat rgb;
    cvtColor(src, rgb, CV_BGR2RGB);

    QImage qTemp = QImage(rgb.data, rgb.cols, rgb.rows, rgb.step, QImage::Format_RGB888);

    if(!qTemp.isNull()) qTemp = qTemp.scaled(ui->iconLabel->size(), Qt::KeepAspectRatio);
    else return;

    QPixmap pix = QPixmap::fromImage(qTemp, 0);
    ui->iconLabel->setPixmap(pix);

    QPixmap *p = (QPixmap *) ui->iconLabel->pixmap();
    QPainter pp;
    pp.begin(p);
    pp.drawRect(x, y, w, h);
    repaint();
    pp.end();
}

/*
 * Draw grid of size 'grid' on pixmap
 */
void TrainingControl::drawLine(int grid)
{
    if(!dst.data) return;

    Mat rgb;
    cvtColor(dst, rgb, CV_BGR2RGB);

    QImage qTemp = QImage(rgb.data, rgb.cols, rgb.rows, rgb.step, QImage::Format_RGB888);
    qTemp = qTemp.scaled(ui->imgLabel->size(), Qt::KeepAspectRatio);

    QPixmap pix;

    if(!qTemp.isNull()) pix = QPixmap::fromImage(qTemp, 0);
    else return;

    QPainter painter(&pix);

    // Top left corner
    int x1 = pix.rect().topLeft().x();
    int y1 = pix.rect().topLeft().y();

    // Bottom right corner
    int x2 = pix.rect().bottomRight().x();
    int y2 = pix.rect().bottomRight().y();

    // Draw horizontal lines
    for(int i = y1; i < y2; i += grid)
    {
        painter.drawLine(x1, i, x2, i);
    }

    // Draw vertical lines
    for(int j = x1; j  < x2; j += grid)
    {
        painter.drawLine(j, y1, j, y2);
    }

    ui->imgLabel->setPixmap(pix);
}

void TrainingControl::mousePressEvent(QMouseEvent *e)
{
    int dx = e->x() - ui->imgLabel->x();
    int dy = e->y() - ui->imgLabel->y();

    if(dx>=0 && dx<=ui->imgLabel->width() && dy>=0 && dy<=ui->imgLabel->height())
    {
        origin = e->pos();
        rubberBand->setGeometry(QRect(origin, QSize()));
        rubberBand->show();
    }
}

void TrainingControl::mouseMoveEvent(QMouseEvent *e)
{

    rubberBand->setGeometry(QRect(origin, e->pos()).normalized());
}

void TrainingControl::mouseReleaseEvent(QMouseEvent *e)
{
    if(img.data)
    {
        // Hide displayed rubber band
        rubberBand->hide();

        int diff_w = e->x() - origin.x();
        int diff_h = e->y() - origin.y();

        int topx = ui->imgLabel->x();
        int topy = ui->imgLabel->y();

        if(diff_w > 0 && diff_w < ui->imgLabel->width() &&
                diff_h > 0 && diff_h < ui->imgLabel->height())
        {
            // Set ROI to both img and resultImg. ROI is
            // manually drawn by user. Although it will not
            // shown
            setROI(origin.x()-topx, origin.y()-topy, diff_w, diff_h);
            // Locate seeds in img and get result image
            // with ROI set for both images
            // gt->manualContour(img, resultImg, runNNFlag);

            // When 'Train Neural Network' Button is click,
            // 'runNNFlag' will be set to true
            if(runNNFlag)
            {
                gt->runNeuralNetwork(img);
                runNNFlag = false;
            }

            Mat rgb;
            cvtColor(resultImg, rgb, CV_BGR2RGB);
            // Display result image
            display(resultImg);
            // Display number of identified seeds
            ui->totalLabel->setText(QString::number(gt->getTotalSeeds()));
            ui->sumLabel->setText(QString::number(gt->getSum()));
        }
    }
}

/*
 * Once the ROI is drawn on UI, map the dimension(x, y, w, h)
 * to original size image. Since pixmap is left and top aligned,
 * origins will be at topleft corner (0, 0). Ratio will be size
 * of original image over size of imgLabel pixmap
 */
void TrainingControl::setROI(int x, int y, int w, int h)
{
    int cols = img.cols;
    int rows = img.rows;

    QPixmap *pix = (QPixmap *)ui->imgLabel->pixmap();

    // get current image side
     int l_width = pix->width();
     int l_height = pix->height();


    // Ratio of mapping from pixmap to
    int r = cols / l_width;

    // Remainders of width and height
    int rw = cols % l_width;
    int rh = rows % l_height;

    QPainter painter;
    painter.begin(pix);
    Rect roi = Rect(x*r, y*r, w*r+rw, h*r+rh);
    painter.drawRect(QRect(x*r, y*r, w*r+rw, h*r+rh));
    painter.end();


    img = Mat(img, roi);
    resultImg = Mat(resultImg, roi);

    displayIcon(dst, roi);
}

void TrainingControl::autoD()
{
    if(img.data)
    {
        gt->autoContour(img, resultImg, runNNFlag);
        if(runNNFlag)
        {
            gt->runNeuralNetwork(img);
            runNNFlag = false;
        }
        Mat rgb;
        cvtColor(resultImg, rgb, CV_BGR2RGB);
        // Display result image
        display(resultImg);
        // Display number of identified seeds
        ui->totalLabel->setText(QString::number(gt->getTotalSeeds()));
        ui->sumLabel->setText(QString::number(gt->getSum()));
        ui->autoButton->setEnabled(false);
    }
}

void TrainingControl::manualD()
{
    QMessageBox msgBox;
    msgBox.setText("Please select image ROI by drawing a rectangle starting "
                   "from top left corner to bottom right");
    msgBox.exec();

    ///////////////////////////////////////////////////////////////
    // Waiting for user to draw rectangle on the image to indicate
    // region of interest
    ///////////////////////////////////////////////////////////////

    drawLine(32);
}

void TrainingControl::setGrid(QString size)
{
    int grid = size.toInt(0, 10);

    drawLine(grid);
}

/*
 *  Reset image label with only original image displayed
 */
void TrainingControl::reset()
{
    display(img);
    ui->iconLabel->setPixmap(QPixmap());
    ui->autoButton->setEnabled(true);
}

/*
 *  Cancel and close this window
 */
void TrainingControl::cancel()
{
    this->close();
}

